import os.path
import pandas as pd
import time
from main import text_to_sql
from project_path import get_project_path
from utils.question_match import load_common_qa


def sql_evaluator(question, answer) -> tuple:
    """
    评估生成 SQL 的准确性
    :param question: 测试问题
    :param answer: 标准答案
    :return: 评估结果(包括生成 SQL,响应耗时)
    """
    is_match = "Mismatch"  # 是否匹配
    start_time = time.time()
    text2sql_answer = text_to_sql(question)
    run_time = round(time.time() - start_time, 2)
    if text2sql_answer == answer:
        is_match = "Match"
    return text2sql_answer, run_time, is_match


def text2sql_test(df: pd.DataFrame):
    """
    text2sql-mcp-server text-to-sql 性能测试
    :param df: 测试数据
    """
    # 测试问题
    questions = df['qa_example'].tolist()
    # 测试答案
    answers = df['qa_sql'].tolist()
    # 记录测试结果
    for question, answer in zip(questions, answers):
        text2sql_answer, run_time, is_match = sql_evaluator(question, answer)
        print("Question: {}\nExpected Answer: {}\nGenerate Answer: {}\nRun: {} seconds\n{}!!!\n"
              .format(question, answer, text2sql_answer, run_time, is_match))
        df.loc[df['qa_example'] == question, 'Generate Answer'] = text2sql_answer
        df.loc[df['qa_example'] == question, 'Run Time'] = run_time
        df.loc[df['qa_example'] == question, 'is Match'] = is_match
    # 保存评估结果
    df.rename(columns={'qa_example': 'Question', 'qa_sql': 'Expected Answer'}, inplace=True)
    df[['Question', 'Expected Answer', 'Generate Answer', 'Run Time', 'Is Match']].to_csv(
        os.path.join(get_project_path(), 'data', 'evaluation_result.csv'), index=False)


if __name__ == '__main__':
    qa_df, qa_mapping, layer_dict = load_common_qa()
    text2sql_test(qa_df)
